"""
生成合成MHD数据示例
"""

import os
import sys
import argparse
import h5py
import numpy as np
import torch

# 添加项目根目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from em_generation.data.em_dataset import SyntheticEMDataGenerator
from em_generation.config import get_config


def main():
    parser = argparse.ArgumentParser(description='生成合成MHD数据')
    parser.add_argument('--output_dir', type=str, default='./data/synthetic', 
                        help='输出目录')
    parser.add_argument('--num_samples', type=int, default=100, 
                        help='样本数量')
    parser.add_argument('--grid_size', type=int, default=32, 
                        help='网格尺寸')
    parser.add_argument('--time_steps', type=int, default=10, 
                        help='时间步数')
    parser.add_argument('--seed', type=int, default=42, 
                        help='随机种子')
    
    args = parser.parse_args()
    
    # 创建输出目录
    os.makedirs(args.output_dir, exist_ok=True)
    
    # 设置随机种子
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    
    # 获取配置
    config = get_config('mhd')
    config.spatial_dim = [args.grid_size, args.grid_size, args.grid_size]
    config.time_steps = args.time_steps
    
    # 创建数据生成器
    generator = SyntheticEMDataGenerator(config)
    
    # 生成数据
    print(f"生成 {args.num_samples} 个MHD样本...")
    generator.generate_mhd_data(
        output_path=os.path.join(args.output_dir, 'mhd_data.h5'),
        num_samples=args.num_samples,
        grid_size=args.grid_size,
        time_steps=args.time_steps
    )
    
    print(f"数据已保存到 {args.output_dir}")


if __name__ == "__main__":
    main()